feat(iorails)!: Support non-streaming GenerationResponse - #2178
Conversation
github-actions
Bot
added
size: L
status: needs triage
Codecov Report✅ All modified and coverable lines are covered by tests. 📢 Thoughts on this report? Let us know! |
tgasser-nv
changed the title
tgasser-nv
added
status: triaged
…d_actions.llm_calls.{prompt, completion}
tgasser-nv
changed the title
|
@coderabbitai Review this PR |
|
@greptile-apps Review this PR |
|
✅ Action performedReview finished.
|
📝 WalkthroughWalkthroughGeneration APIs now accept typed options. IORails can return structured ChangesStructured generation and rail telemetry
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client
participant IORails
participant RailsManager
participant RailAction
participant EngineRegistry
participant GenerationLog
Client->>IORails: generate(options)
IORails->>RailsManager: run safety rails
RailsManager->>RailAction: execute rail
RailAction->>EngineRegistry: model/API call
EngineRegistry-->>RailAction: response and call metadata
RailsManager-->>IORails: RailResult with RailCallRecord
IORails->>EngineRegistry: main model generation
IORails->>GenerationLog: aggregate records and stats
IORails-->>Client: GenerationResponse
Possibly related PRs
🚥 Pre-merge checks | ✅ 6✅ Passed checks (6 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
|
@coderabbitai Review this PR |
|
✅ Action performedReview finished.
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
nemoguardrails/guardrails/iorails.py (1)
958-1067: 🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy liftCapture speculative generation before applying the input-rail verdict.
When generation finishes first and input rails later block, Lines 1051-1058 return without appending the completed main call at Line 1067. Simultaneous completion in the rails-first branch has the same omission. This undercounts calls, tokens, and cost on blocked requests; successful speculative calls also lose timing. Time and record the generation task at its execution boundary, regardless of whether its response is ultimately returned.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@nemoguardrails/guardrails/iorails.py` around lines 958 - 1067, Update _parallel_input_rail_and_response_generation to capture and record the completed generation task before applying the input-rail verdict, including both generation-first and simultaneous-completion paths. Time the main LLM call at its execution boundary and append its generation record with usage, model, timestamps, duration, provider, and main_prompt even when input rails later reject the response; avoid recording the call twice.nemoguardrails/guardrails/rails_manager.py (1)
438-461: 🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick winProcess every completed task before returning an unsafe result.
asyncio.waitcan place multiple tasks indone. If the first sorted result is unsafe, Line 460 returns before collecting records—or retrieving exceptions—from the remaining completed tasks. This violates the “every rail that ran” log contract. Consume the fulldonebatch, select the first unsafe result, then cancel only pending tasks.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@nemoguardrails/guardrails/rails_manager.py` around lines 438 - 461, Update the task-processing loop in the rail execution method to process every task in each `done` batch, collecting records and retrieving results or exceptions before deciding to return. Track the first unsafe result encountered in sorted task order, then cancel and await only `pending_tasks`; return that unsafe result with all collected records after the batch is fully consumed.
🧹 Nitpick comments (1)
tests/guardrails/test_guardrails.py (1)
126-127: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd a non-
Noneoptions forwarding test.These assertions only cover the default path. Add one sync/async wrapper test with a sentinel
GenerationOptionsor dict and assert that the same object is passed to the underlying engine, preventing regressions that accept but drop the new argument.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/guardrails/test_guardrails.py` around lines 126 - 127, Add a test covering non-None options for both sync and async guardrail wrappers, using a sentinel GenerationOptions instance or dict. Invoke the wrappers with that same object and assert rails_engine.generate and generate_async each receive it unchanged alongside messages, while preserving the existing default-path assertions.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@nemoguardrails/guardrails/guardrails_types.py`:
- Around line 134-141: Update serialize_prompt so captured prompts preserve all
message metadata, including tool_calls, tool_call_id, names, and reasoning-only
fields, instead of formatting only role and content. Reuse the existing
canonical message serialization if available; otherwise encode each complete
message structure while retaining the established prompt ordering and
separation. Add coverage for tool-call and reasoning-only turns.
In `@nemoguardrails/guardrails/guardrails.py`:
- Around line 221-255: The generate_async overloads currently share identical
parameters, causing type checkers to select only the first str return type.
Update generate_async so overloads distinguish calls by the options value, or
replace them with a single accurate union return annotation covering str, dict,
GenerationResponse, and tuple[dict, dict].
In `@nemoguardrails/guardrails/rail_action.py`:
- Around line 197-221: Update both LLM-call helpers surrounding
EngineRegistry.model_call and the corresponding helper at the referenced later
section to record attempted calls when the await raises. In each exception path,
capture elapsed timing and populate _rail_llm_call_var with the available
request/model/provider and failure-safe metadata, then re-raise the original
exception so existing error handling remains unchanged. Preserve the current
successful-response recording behavior.
---
Outside diff comments:
In `@nemoguardrails/guardrails/iorails.py`:
- Around line 958-1067: Update _parallel_input_rail_and_response_generation to
capture and record the completed generation task before applying the input-rail
verdict, including both generation-first and simultaneous-completion paths. Time
the main LLM call at its execution boundary and append its generation record
with usage, model, timestamps, duration, provider, and main_prompt even when
input rails later reject the response; avoid recording the call twice.
In `@nemoguardrails/guardrails/rails_manager.py`:
- Around line 438-461: Update the task-processing loop in the rail execution
method to process every task in each `done` batch, collecting records and
retrieving results or exceptions before deciding to return. Track the first
unsafe result encountered in sorted task order, then cancel and await only
`pending_tasks`; return that unsafe result with all collected records after the
batch is fully consumed.
---
Nitpick comments:
In `@tests/guardrails/test_guardrails.py`:
- Around line 126-127: Add a test covering non-None options for both sync and
async guardrail wrappers, using a sentinel GenerationOptions instance or dict.
Invoke the wrappers with that same object and assert rails_engine.generate and
generate_async each receive it unchanged alongside messages, while preserving
the existing default-path assertions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 0f428c31-8892-4e50-a8d1-5398aa3ee1be
📒 Files selected for processing (19)
nemoguardrails/base_guardrails.pynemoguardrails/guardrails/actions/content_safety_action.pynemoguardrails/guardrails/actions/jailbreak_detection_action.pynemoguardrails/guardrails/actions/topic_safety_action.pynemoguardrails/guardrails/engine_registry.pynemoguardrails/guardrails/guardrails.pynemoguardrails/guardrails/guardrails_types.pynemoguardrails/guardrails/iorails.pynemoguardrails/guardrails/rail_action.pynemoguardrails/guardrails/rails_manager.pytests/guardrails/test_guardrails.pytests/guardrails/test_iorails.pytests/guardrails/test_iorails_generation_log.pytests/guardrails/test_iorails_generation_log_capture.pytests/guardrails/test_iorails_generation_response.pytests/guardrails/test_jailbreak_detection_iorails_actions.pytests/guardrails/test_rails_manager.pytests/guardrails/test_tool_rails_iorails.pytests/guardrails/test_topic_safety_iorails_actions.py
This comment was marked as outdated.
This comment was marked as outdated.
…ith LLMRails signature
|
@greptile-apps Re this comment:
The GenerationOptions / GenerationResponse support for the streaming path will be added in a stacked PR after this one. Check the title of the PR: "feat(iorails)!: Support non-streaming GenerationResponse". Re-evaluate the overall comment and rating |
|
Fair point — the PR title is explicit ("Support non-streaming GenerationResponse") and streaming support is a stated follow-up. I was treating this as an oversight rather than an intentional deferral. That said, there is still a narrow interim risk worth noting: if If the former is the case, a one-liner |
Pouyanpi
left a comment
Pouyanpi
left a comment
There was a problem hiding this comment.
LGTM! A lot of the plumbing here (the call-capture contextvar, the parallel log structs, etc) is really working around IORails not yet having the canonical chat model + http subsystem once the rail-library work lands and rails go through shared infra, most of this should get a lot smoother.
2 participants
Description
Adds the
GenerationResponsereturn argument to non-streaminggenerate()andgenerate_sync()methods. When aGenerationOptionsobject is provided, theGenerationResponsestructured response is returned rather than a dict withrole/content/tool_callskeys.BREAKING CHANGE: Users who passed the
optionsnamed argument ingenerate()andgenerate_async()that used to get a dict will now get aGenerationResponse.BREAKING CHANGE: This PR contains changes to the Guardrails and IORails signatures for these two functions to align them exactly with LLMRails. The signature of
generate()andgenerate_async()now includespromptas the first parameter, to match LLMRails. It changed from:to
This means users who passed a
messagesparameter in the first positional argument ofgenerate_async()orgenerate()will now get aTypeError("prompt must be a string; pass a message list via messages="). Thestateandstreaming_handlerarguments LLMRails supports for these functions aren't supported in IORails.The
GenerationResponsecollects many different items of metadata, timestamps and durations into one place. TheGenerationResponse.logfield has a typeGenerationLog, which in IORails does not support theinternal_eventsandcolang_historyfields. TheGenerationLogOptionsclass controls which fields inGenerationLogare returned. IfGenerationLogOptions.internal_eventsorGenerationLogOptions.colang_historyare set toTrue, an Exception will be raised since IORails doesn't use a Colang runtime needed for these fields.Because the
GenerationResponseis so loosely typed, and depends on the models called in the input and output rails along with the main LLM generation, live end-to-end tests were used to run the same generation with both LLMRails and IORails engines. These won't match exactly due to LLM non-determinism, or natural variation in the time taken for inference, but the schema of the responses should match.GenerationResponsedetailed field support in IORailsFully supported
responsemessages, so always a 1-element[assistant_msg]list.tool_callsToolCall.to_dict()dict-args shape, matching LLMRails'GenerationResponse.tool_calls. (Bare/streaming paths keep the JSON-string wire shape, unchanged.)reasoning_contentreasoningfield, or<think>extracted from the completion; the structured path keepsresponsecontent clean (no inline<think>).llm_outputNone— identical to LLMRails, whose source for this field is never populated either.llm_output=Trueis accepted as a silent no-op.Partially supported
logactivated_rails(one syntheticExecutedAction+LLMCallInfoper rail), the flatllm_callslist, and aggregatestats— all synthesized from per-railRailCallRecords + the main call. Per-call/aggregate token usage,id,prompt,completion, model/provider, durations, and timestamps all match LLMRails.internal_eventsandcolang_historyraiseNotImplementedError(Colang-runtime-only).activated_rails[].decisionsis always[](a Colang construct). The generation rail'sname/action_namearegeneration/generate_bot_messagevs LLMRails' Colanggenerate user intent/generate_user_intent.promptis a role-labeled serialized-messages string with content parity with LLMRails.llm_metadataprovider_metadataverbatim (e.g.response_headers).usagesub-key — all token usage lives inlog(LLMRails mirror). OftenNonewhen the provider returns no metadata blob.Not supported
output_dataNone; passingoptions.output_varsraisesValueError.stateNone; passingoptions.stateraisesValueError.Integration-testing (using
compare_generate.py)This script issues the identical request to both IORails and LLMRails, with the
GenerationLogOptionsfields that IORails doesn't support disabled.It was run as below from the Guardrails uv environment
This generated the
llmrails_generation_response.jsonandiorails_generation_response.jsonfiles attached to the PR.iorails_generation_response.json
llmrails_generation_response.json
The Field-by-field Comparison output is below. For the full log file see here
generation_compare.log .
Related Issue(s)
Verification
Pre-commit
Unit-test
Integration test with Chat
AI Assistance
Checklist
Summary by CodeRabbit